import numpy as np
import torch.optim
import tqdm
from torchvision.datasets import mnist
from torchvision import transforms
from torch.utils.data import DataLoader
from network import Generator,Discriminator

MaxEpoch = 100
rank = 'cuda'
batch_size =32
z_dim = 100


data_tf = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize([0.5], [0.5])])
#------------------------------------#
#   loading MNIST dataset
#------------------------------------#
train_set = mnist.MNIST('./MNIST', train=True, transform=data_tf, download=False)
test_set = mnist.MNIST('./MNIST', train=False, transform=data_tf, download=False)
trainDataLoader = DataLoader(train_set,batch_size=batch_size,shuffle=True,num_workers=4)
testDataLoader = DataLoader(test_set,batch_size=batch_size,shuffle=True,num_workers=4)

#------------------------------------#
#       G
#------------------------------------#
G = Generator(z_dim=z_dim)


#------------------------------------#
#       D
#------------------------------------#
D = Discriminator()
#------------------------------------#
#       Loss Function
#------------------------------------#
loss = torch.nn.BCELoss()

#------------------------------------#
#       Optimizer
#------------------------------------#
G_optim = torch.optim.Adam(G.parameters(),lr=0.0002)
D_optim = torch.optim.Adam(D.parameters(),lr=0.0002)
if rank == 'cuda':
    device = torch.device('cuda')
    G = G.to(device)
    D = D.to(device)
    loss = loss.to(device)
    ones = torch.ones(batch_size).to(device)
    zeros = torch.zeros(batch_size).to(device)

for epoch in range(MaxEpoch):
    G.train()
    D.train()
    with tqdm.tqdm(total=len(trainDataLoader),desc=f'{epoch+1}/{MaxEpoch}:') as bar :
        for data,label in trainDataLoader:
            tmp = torch.FloatTensor(label.size(0), 10).zero_()
            label = tmp.scatter_(dim=1, index=torch.LongTensor(label.view(-1, 1)), value=1)


            z = torch.randn([batch_size, z_dim])
            c = torch.FloatTensor(batch_size, 10).zero_()
            c = c.scatter_(dim=1, index=torch.LongTensor(np.random.choice(10, batch_size).reshape([batch_size, 1])),value=1)
            if rank =='cuda':
                device = torch.device('cuda')
                data = data.to(device)
                label  = label.to(device)
                z = z.to(device)
                c = c.to(device)
            fake_image =G(z,c)

            real_loss = loss(D(data,label),ones.view(-1,1))
            fake_loss = loss(D(fake_image.detach(),c),zeros.view(-1,1))
            d_loss = real_loss + fake_loss
            D_optim.zero_grad()
            d_loss.backward()
            D_optim.step()

            g_loss = loss(D(fake_image,c),ones.view(-1,1))
            G_optim.zero_grad()
            g_loss.backward()
            G_optim.step()
            bar.set_postfix({'G_loss:':d_loss.item(),'D_loss:':d_loss.item()})
            bar.update(1)
    if ((epoch+1) % 40) == 0 :
        torch.save(G.state_dict(),'./weight/cGan.pth')



